{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TransE 实践\n",
    "\n",
    "在这个演示中，我们使用TransE([论文链接](https://www.utc.fr/~bordesan/dokuwiki/_media/en/transe_nips13.pdf))对示例中文知识图谱进行链接预测，从而达到补全知识图谱的目的。\n",
    "\n",
    "希望在这个demo中帮助大家了解知识图谱表示学习的作用原理和机制。\n",
    "\n",
    "本demo建议使用python3运行。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数据集\n",
    "这个示例中，我们使用的是一个人工构造的小的中文数据集，示例图谱可视化如下：\n",
    "<img src=\"./pic/示例图谱可视化.png\" style=\"zoom:50%;\" />\n",
    "\n",
    "train.txt：包含28个训练三元组，文件的每一行每一行表示一个三元组, 按头实体、关系、尾实体顺序 ，并用’\\t’分隔。\n",
    "\n",
    "test.txt：包含4个测试三元组，文件的每一行表示一个三元组, 按头实体、关系、尾实体顺序 ，并用’\\t’分隔。\n",
    "\n",
    "entity2id.txt：存储了实体对应id的信息，示例图谱中共包含17个实体，为每个实体分配一个id从0-16，文件每一行表示一个实体和对应的id并用’\\t’分隔。\n",
    "\n",
    "relation2id.txt： 存储了关系对应id的信息，示例图谱中共包含6个关系，为每个关系分配一个id从0-5，文件每一行表示一个关系和对应的id并用’\\t’分隔\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TransE 原理回顾\n",
    "TransE将每个实体和关系都表示成一个向量，并假设对于一个存在在知识图谱中的三元组$(h,r,t)$, $h, r, t$的向量表示$\\mathbf{h}, \\mathbf{r}, \\mathbf{t}$满足：\n",
    "<img src=\"./pic/TransE-向量空间假设.png\" style=\"zoom:50%;\" />\n",
    "即\n",
    "$$\\mathbf{h} + \\mathbf{r} = \\mathbf{t}$$\n",
    "\n",
    "\n",
    "\n",
    "对于每个正确的三元组的优化目标是：\n",
    "$$\\mathbf{h} + \\mathbf{r} \\approx \\mathbf{t}$$\n",
    "对于一个三元组的评分函数为：\n",
    "$$f_r(h,t) = \\| \\mathbf{h} + \\mathbf{r} - \\mathbf{t} \\| _{L_1/L_2} $$\n",
    "TransE的损失函数：\n",
    "$$ L = \\sum_{(h,r,t)\\in S} \\sum_{(h^\\prime, r^\\prime, t^\\prime) \\in S^\\prime} max(0, f_r(h,t) + \\gamma - f_{r^\\prime} (h^\\prime, t^\\prime)) $$\n",
    "其中$S$是所有正样本的集合，$S^\\prime$是所有负样本的集合，对于一个正样本$(h,r,t)$负样本通过随机替换$h$或$t$得到， $\\gamma$表示间隔，是一个超参。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 代码实践"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf \n",
    "import time \n",
    "import argparse\n",
    "import random\n",
    "import numpy as np \n",
    "import os.path\n",
    "import math\n",
    "import timeit\n",
    "from multiprocessing import JoinableQueue, Queue, Process\n",
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransE:\n",
    "    @property\n",
    "    def variables(self):\n",
    "        return self.__variables\n",
    "\n",
    "    @property\n",
    "    def num_triple_train(self):\n",
    "        return self.__num_triple_train\n",
    "\n",
    "    @property \n",
    "    def num_triple_test(self):\n",
    "        return self.__num_triple_test\n",
    "\n",
    "    @property\n",
    "    def testing_data(self):\n",
    "        return self.__triple_test\n",
    "\n",
    "    @property \n",
    "    def num_entity(self):\n",
    "        return self.__num_entity\n",
    "\n",
    "    @property\n",
    "    def embedding_entity(self):\n",
    "        return self.__embedding_entity\n",
    "\n",
    "\n",
    "    @property\n",
    "    def embedding_relation(self):\n",
    "        return self.__embedding_relation\n",
    "\n",
    "    @property\n",
    "    def hr_t(self):\n",
    "        return self.__hr_t\n",
    "\n",
    "    @property \n",
    "    def tr_h(self):\n",
    "        return self.__tr_h\n",
    "    \n",
    "    @property\n",
    "    def entity2id(self):\n",
    "        return self.__entity2id\n",
    "    \n",
    "    @property\n",
    "    def relation2id(self):\n",
    "        return self.__relation2id\n",
    "\n",
    "    @property\n",
    "    def id2entity(self):\n",
    "        return self.__id2entity\n",
    "    \n",
    "    @property\n",
    "    def id2relation(self):\n",
    "        return self.__id2relation\n",
    "\n",
    "    def training_data_batch(self, batch_size = 512):\n",
    "        n_triple = len(self.__triple_train)\n",
    "        rand_idx = np.random.permutation(n_triple)\n",
    "        start = 0\n",
    "        while start < n_triple:\n",
    "            start_t = timeit.default_timer()\n",
    "            end = min(start+batch_size, n_triple)\n",
    "            size = end - start \n",
    "            train_triple_positive = np.asarray([ self.__triple_train[x] for x in  rand_idx[start:end]])\n",
    "            train_triple_negative = []\n",
    "            num_negative = 5 \n",
    "            for t in train_triple_positive:\n",
    "                replace_entity_id = np.random.randint(self.__num_entity)\n",
    "                random_num = np.random.random()\n",
    "\n",
    "                if self.__negative_sampling == 'unif':\n",
    "                    replace_head_probability = 0.5\n",
    "                elif self.__negative_sampling == 'bern':\n",
    "                    replace_head_probability = self.__relation_property[t[1]]\n",
    "                else:\n",
    "                    raise NotImplementedError(\"Dose not support %s negative_sampling\" %negative_sampling)\n",
    "\n",
    "                if random_num<replace_head_probability:\n",
    "                    train_triple_negative.append((replace_entity_id, t[1],t[2]))\n",
    "                else:\n",
    "                    train_triple_negative.append((t[0], t[1], replace_entity_id))\n",
    "\n",
    "            start = end\n",
    "            prepare_t = timeit.default_timer()-start_t\n",
    "\n",
    "            yield train_triple_positive, train_triple_negative, prepare_t\n",
    "\n",
    "\n",
    "    def __init__(self, data_dir, negative_sampling,learning_rate, \n",
    "             batch_size, max_iter, margin, dimension, norm, evaluation_size, regularizer_weight):\n",
    "        # this part for data prepare\n",
    "        self.__data_dir=data_dir\n",
    "        self.__negative_sampling=negative_sampling\n",
    "        self.__regularizer_weight = regularizer_weight\n",
    "        self.__norm = norm\n",
    "\n",
    "        self.__entity2id={}\n",
    "        self.__id2entity={}\n",
    "        self.__relation2id={}\n",
    "        self.__id2relation={}\n",
    "\n",
    "        self.__triple_train=[] #[(head_id, relation_id, tail_id),...]\n",
    "        self.__triple_test=[]\n",
    "        self.__triple_valid=[]\n",
    "        self.__triple = []\n",
    "\n",
    "        self.__num_entity=0\n",
    "        self.__num_relation=0\n",
    "        self.__num_triple_train=0\n",
    "        self.__num_triple_test=0\n",
    "        self.__num_triple_valid=0\n",
    "\n",
    "        # load all the file: entity2id.txt, relation2id.txt, train.txt, test.txt, valid.txt\n",
    "        self.load_data()\n",
    "        print('finish preparing data. ')\n",
    "\n",
    "\n",
    "        # this part for the model:\n",
    "        self.__learning_rate = learning_rate\n",
    "        self.__batch_size = batch_size\n",
    "        self.__max_iter = max_iter\n",
    "        self.__margin = margin\n",
    "        self.__dimension = dimension\n",
    "        self.__variables= []\n",
    "        #self.__norm = norm\n",
    "        self.__evaluation_size = evaluation_size\n",
    "        bound = 6 / math.sqrt(self.__dimension)\n",
    "        with tf.device('/cpu'):\n",
    "            self.__embedding_entity = tf.get_variable('embedding_entity', [self.__num_entity, self.__dimension],\n",
    "                                                       initializer=tf.random_uniform_initializer(minval=-bound, maxval=bound, seed = 123))\n",
    "            self.__embedding_relation = tf.get_variable('embedding_relation', [self.__num_relation, self.__dimension],\n",
    "                                                         initializer=tf.random_uniform_initializer(minval=-bound, maxval=bound, seed =124))\n",
    "            self.__variables.append(self.__embedding_entity)\n",
    "            self.__variables.append(self.__embedding_relation)\n",
    "            print('finishing initializing')\n",
    "\n",
    "\n",
    "    def load_data(self):\n",
    "        print('loading entity2id.txt ...')\n",
    "        with open(os.path.join(self.__data_dir, 'entity2id.txt'), encoding='utf-8') as f:\n",
    "            self.__entity2id = {line.strip().split('\\t')[0]: int(line.strip().split('\\t')[1]) for line in f.readlines()}\n",
    "            self.__id2entity = {value:key for key,value in self.__entity2id.items()}\n",
    "\n",
    "        print('loading reltion2id.txt ...')     \n",
    "        with open(os.path.join(self.__data_dir,'relation2id.txt'), encoding='utf-8') as f:\n",
    "            self.__relation2id = {line.strip().split('\\t')[0]: int(line.strip().split('\\t')[1]) for line in f.readlines()}\n",
    "            self.__id2relation = {value:key for key, value in self.__relation2id.items()}\n",
    "\n",
    "        def load_triple(self, triplefile):\n",
    "            triple_list = [] #[(head_id, relation_id, tail_id),...]\n",
    "            with open(os.path.join(self.__data_dir, triplefile), encoding='utf-8') as f:\n",
    "                for line in f.readlines():\n",
    "                    line_list = line.strip().split('\\t')\n",
    "                    assert len(line_list) == 3\n",
    "                    headid = self.__entity2id[line_list[0]]\n",
    "                    relationid = self.__relation2id[line_list[1]]\n",
    "                    tailid = self.__entity2id[line_list[2]]\n",
    "                    triple_list.append((headid, relationid, tailid))\n",
    "                    self.__hr_t[(headid, relationid)].add(tailid)\n",
    "                    self.__tr_h[(tailid, relationid)].add(headid)\n",
    "            return triple_list\n",
    "\n",
    "        self.__hr_t = defaultdict(set)\n",
    "        self.__tr_h = defaultdict(set)\n",
    "        self.__triple_train = load_triple(self, 'train.txt')\n",
    "        self.__triple_test = load_triple(self, 'test.txt')\n",
    "        self.__triple_valid = load_triple(self, 'valid.txt')\n",
    "        self.__triple = np.concatenate([self.__triple_train, self.__triple_test, self.__triple_valid], axis = 0 )\n",
    "\n",
    "        self.__num_relation = len(self.__relation2id)\n",
    "        self.__num_entity = len(self.__entity2id)\n",
    "        self.__num_triple_train = len(self.__triple_train)\n",
    "        self.__num_triple_test = len(self.__triple_test)\n",
    "        self.__num_triple_valid = len(self.__triple_valid)\n",
    "\n",
    "        print('entity number: ' + str(self.__num_entity))\n",
    "        print('relation number: ' + str(self.__num_relation))\n",
    "        print('training triple number: ' + str(self.__num_triple_train))\n",
    "        print('testing triple number: ' + str(self.__num_triple_test))\n",
    "        print('valid triple number: ' + str(self.__num_triple_valid))\n",
    "\n",
    "\n",
    "        if self.__negative_sampling == 'bern':\n",
    "            self.__relation_property_head = {x:[] for x in range(self.__num_relation)} #{relation_id:[headid1, headid2,...]}\n",
    "            self.__relation_property_tail = {x:[] for x in range(self.__num_relation)} #{relation_id:[tailid1, tailid2,...]}\n",
    "            self.__relation_property = {x:[] for x in range(self.__num_relation)} \n",
    "            for t in self.__triple_train:\n",
    "                #print(t)\n",
    "                self.__relation_property_head[t[1]].append(t[0])\n",
    "                self.__relation_property_tail[t[1]].append(t[2])\n",
    "            #print(self.__relation_property_head[0])\n",
    "            #print(self.__relation_property_tail[0])\n",
    "            for x in self.__relation_property_head.keys():\n",
    "                t = len(set(self.__relation_property_tail[x]))\n",
    "                h = len(set(self.__relation_property_head[x]))\n",
    "                self.__relation_property[x] = float(t)/(h+t+0.000000001)\n",
    "            #self.__relation_property = {x:(len(set(self.__relation_property_tail[x])))/(len(set(self.__relation_property_head[x]))+ len(set(self.__relation_property_tail[x]))) \\\n",
    "            #\t\t\t\t\t\t\t for x in self.__relation_property_head.keys()} # {relation_id: p, ...} 0< num <1, and for relation replace head entity with the property p\n",
    "        else: \n",
    "            print(\"unif set don't need to calculate hpt and tph\")\n",
    "\n",
    "\n",
    "\n",
    "    def train(self, inputs):\n",
    "        embedding_relation = self.__embedding_relation\n",
    "        embedding_entity = self.__embedding_entity\n",
    "\n",
    "        triple_positive, triple_negative = inputs # triple_positive:(head_id,relation_id,tail_id)\n",
    "\n",
    "        norm_entity = tf.nn.l2_normalize(embedding_entity, dim = 1)\n",
    "        norm_relation = tf.nn.l2_normalize(embedding_relation, dim = 1)\n",
    "        norm_entity_l2sum = tf.sqrt(tf.reduce_sum(norm_entity**2, axis = 1))\n",
    "\n",
    "        embedding_positive_head = tf.nn.embedding_lookup(norm_entity, triple_positive[:, 0])\n",
    "        embedding_positive_tail = tf.nn.embedding_lookup(norm_entity, triple_positive[:, 2])\n",
    "        embedding_positive_relation = tf.nn.embedding_lookup(norm_relation, triple_positive[:, 1])\n",
    "\n",
    "        embedding_negative_head = tf.nn.embedding_lookup(norm_entity, triple_negative[:, 0])\n",
    "        embedding_negative_tail = tf.nn.embedding_lookup(norm_entity, triple_negative[:, 2])\n",
    "        embedding_negative_relation = tf.nn.embedding_lookup(norm_relation, triple_negative[:, 1])\n",
    "\n",
    "        score_positive = tf.reduce_sum(tf.abs(embedding_positive_head + embedding_positive_relation - embedding_positive_tail), axis = 1)\n",
    "        score_negative = tf.reduce_sum(tf.abs(embedding_negative_head + embedding_negative_relation - embedding_negative_tail), axis = 1)\n",
    "\n",
    "        loss_every = tf.maximum(0., score_positive + self.__margin - score_negative)\n",
    "        loss_triple = tf.reduce_sum(tf.maximum(0., score_positive + self.__margin - score_negative))\n",
    "        self.__loss_regularizer = loss_regularizer = tf.reduce_sum(tf.abs(self.__embedding_relation)) + tf.reduce_sum(tf.abs(self.__embedding_entity))\n",
    "        return loss_triple, loss_every, norm_entity_l2sum #+ loss_regularizer*self.__regularizer_weight\n",
    "\n",
    "    def test(self, inputs):\n",
    "        embedding_relation = self.__embedding_relation\n",
    "        embedding_entity = self.__embedding_entity\n",
    "\n",
    "        triple_test = inputs # (headid, tailid, tailid)\n",
    "        head_vec = tf.nn.embedding_lookup(embedding_entity, triple_test[0])\n",
    "        rel_vec = tf.nn.embedding_lookup(embedding_relation, triple_test[1])\n",
    "        tail_vec = tf.nn.embedding_lookup(embedding_entity, triple_test[2])\n",
    "\n",
    "        norm_embedding_entity = tf.nn.l2_normalize(embedding_entity, dim =1 )\n",
    "        norm_embedding_relation = tf.nn.l2_normalize(embedding_relation, dim = 1)\n",
    "        norm_head_vec = tf.nn.embedding_lookup(norm_embedding_entity, triple_test[0])\n",
    "        norm_rel_vec = tf.nn.embedding_lookup(norm_embedding_relation, triple_test[1])\n",
    "        norm_tail_vec = tf.nn.embedding_lookup(norm_embedding_entity, triple_test[2])\n",
    "\n",
    "        _, id_replace_head = tf.nn.top_k(tf.reduce_sum(tf.abs(embedding_entity + rel_vec - tail_vec), axis=1), k=self.__num_entity)\n",
    "        _, id_replace_tail = tf.nn.top_k(tf.reduce_sum(tf.abs(head_vec + rel_vec - embedding_entity), axis=1), k=self.__num_entity)\n",
    "\n",
    "        _, norm_id_replace_head = tf.nn.top_k(tf.reduce_sum(tf.abs(norm_embedding_entity + norm_rel_vec - norm_tail_vec), axis=1), k=self.__num_entity)\n",
    "        _, norm_id_replace_tail = tf.nn.top_k(tf.reduce_sum(tf.abs(norm_head_vec + norm_rel_vec - norm_embedding_entity), axis=1), k=self.__num_entity)\n",
    "\n",
    "        return id_replace_head, id_replace_tail, norm_id_replace_head, norm_id_replace_tail"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_operation(model, learning_rate=0.01, margin=1.0, optimizer_str = 'gradient'):\n",
    "    with tf.device('/cpu'):\n",
    "        train_triple_positive_input = tf.placeholder(tf.int32, [None, 3])\n",
    "        train_triple_negative_input = tf.placeholder(tf.int32, [None, 3])\n",
    "\n",
    "        loss, loss_every, norm_entity = model.train([train_triple_positive_input, train_triple_negative_input])\n",
    "        if optimizer_str == 'gradient':\n",
    "            optimizer = tf.train.GradientDescentOptimizer(learning_rate = learning_rate)\n",
    "        elif optimizer_str == 'rms':\n",
    "            optimizer = tf.train.RMSPropOptimizer(learning_rate = learning_rate)\n",
    "        elif optimizer_str == 'adam':\n",
    "            optimizer = tf.train.AdamOptimizer(learning_rate = learning_rate)\n",
    "        else:\n",
    "            raise NotImplementedError(\"Dose not support %s optimizer\" %optimizer_str)\n",
    "\n",
    "        grads = optimizer.compute_gradients(loss, model.variables)\n",
    "        op_train = optimizer.apply_gradients(grads)\n",
    "\n",
    "        return train_triple_positive_input, train_triple_negative_input, loss, op_train, loss_every, norm_entity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_operation(model):\n",
    "    with tf.device('/cpu'):\n",
    "        test_triple = tf.placeholder(tf.int32, [3])\n",
    "        head_rank, tail_rank, norm_head_rank, norm_tail_rank = model.test(test_triple)\n",
    "        return test_triple, head_rank, tail_rank, norm_head_rank, norm_tail_rank"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_job(inputqueue, outputqueue, hr_t, tr_h):\n",
    "    while True:\n",
    "        dat = inputqueue.get()\n",
    "        if dat is None:\n",
    "            inputqueue.task_done()\n",
    "            continue\n",
    "        t, id_replace_head, id_replace_tail, norm_id_replace_head, norm_id_replace_tail = dat \n",
    "        hrank = 0\n",
    "        fhrank = 0\n",
    "        for i in range(len(id_replace_head)):\n",
    "            val = id_replace_head[-i-1]\n",
    "            if val == t[0]:\n",
    "                break\n",
    "            else: \n",
    "                hrank += 1\n",
    "                fhrank += 1 \n",
    "                if val in tr_h[(t[2],t[1])]:\n",
    "                    fhrank -= 1\n",
    "\n",
    "        norm_hrank = 0\n",
    "        norm_fhrank = 0\n",
    "        for i in range(len(norm_id_replace_head)):\n",
    "            val = norm_id_replace_head[-i-1]\n",
    "            if val == t[0]:\n",
    "                break\n",
    "            else: \n",
    "                norm_hrank += 1\n",
    "                norm_fhrank += 1 \n",
    "                if val in tr_h[(t[2],t[1])]:\n",
    "                    norm_fhrank -= 1\n",
    "\n",
    "\n",
    "        trank = 0\n",
    "        ftrank = 0\n",
    "        for i in range(len(id_replace_tail)):\n",
    "            val = id_replace_tail[-i-1]\n",
    "            if val == t[2]:\n",
    "                break\n",
    "            else:\n",
    "                trank += 1\n",
    "                ftrank += 1\n",
    "                if val in hr_t[(t[0], t[1])]:\n",
    "                    ftrank -= 1\n",
    "\n",
    "        norm_trank = 0\n",
    "        norm_ftrank = 0\n",
    "        for i in range(len(norm_id_replace_tail)):\n",
    "            val = norm_id_replace_tail[-i-1]\n",
    "            if val == t[2]:\n",
    "                break\n",
    "            else:\n",
    "                norm_trank += 1\n",
    "                norm_ftrank += 1\n",
    "                if val in hr_t[(t[0], t[1])]:\n",
    "                    norm_ftrank -= 1\n",
    "        #print(hrank, fhrank, trank, ftrank,norm_hrank, norm_fhrank, norm_trank, norm_ftrank)\n",
    "        outputqueue.put((hrank, fhrank, trank, ftrank,norm_hrank, norm_fhrank, norm_trank, norm_ftrank))\n",
    "        inputqueue.task_done()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Args:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<__main__.Args object at 0xb3cbc8898>\n",
      "loading entity2id.txt ...\n",
      "loading reltion2id.txt ...\n",
      "entity number: 17\n",
      "relation number: 6\n",
      "training triple number: 28\n",
      "testing triple number: 4\n",
      "valid triple number: 4\n",
      "finish preparing data. \n",
      "finishing initializing\n",
      "WARNING:tensorflow:From <ipython-input-2-6c2b1e1aec70>:209: calling l2_normalize (from tensorflow.python.ops.nn_impl) with dim is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "dim is deprecated, use axis instead\n"
     ]
    }
   ],
   "source": [
    "# 设置参数等\n",
    "args  = Args()\n",
    "args.data_dir = './data/'\n",
    "args.learning_rate = 0.005\n",
    "args.batch_size = 1\n",
    "args.max_iter = 100\n",
    "args.optimizer = 'adam'\n",
    "args.dimension = 10\n",
    "args.margin = 1.0\n",
    "args.norm = 'L1'\n",
    "args.evaluation_size = 500\n",
    "args.save_dir = 'output/'\n",
    "args.negative_sampling = 'bern'\n",
    "args.evaluate_per_iteration = 10\n",
    "args.evaluate_worker = 3\n",
    "args.regularizer_weight = 1e-5\n",
    "args.n_test = 1\n",
    "args.save_per = 100\n",
    "args.n_worker = 5\n",
    "print(args)\n",
    "model = TransE(negative_sampling=args.negative_sampling, data_dir=args.data_dir,\n",
    "                learning_rate=args.learning_rate, batch_size=args.batch_size,\n",
    "                max_iter=args.max_iter, margin=args.margin, \n",
    "                dimension=args.dimension, norm=args.norm, evaluation_size=args.evaluation_size, \n",
    "                regularizer_weight = args.regularizer_weight)\n",
    "\n",
    "train_triple_positive_input, train_triple_negative_input, loss, op_train, loss_every, norm_entity = train_operation(model, learning_rate = args.learning_rate, margin = args.margin, optimizer_str = args.optimizer)\n",
    "test_triple, head_rank, tail_rank , norm_head_rank, norm_tail_rank= test_operation(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /Users/zhangwen/anaconda3/envs/tf1.12/lib/python3.6/site-packages/tensorflow/python/util/tf_should_use.py:189: initialize_all_variables (from tensorflow.python.ops.variables) is deprecated and will be removed after 2017-03-02.\n",
      "Instructions for updating:\n",
      "Use `tf.global_variables_initializer` instead.\n",
      "iter[0] ---loss: 31.93908 ---time: 0.66 ---prepare time : 0.00\n",
      "Model saved at ./save/TransE_0.ckpt\n",
      "iter[1] ---loss: 36.24439 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[2] ---loss: 22.75225 ---time: 0.07 ---prepare time : 0.00\n",
      "iter[3] ---loss: 15.13599 ---time: 0.07 ---prepare time : 0.00\n",
      "iter[4] ---loss: 17.93469 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[5] ---loss: 14.20885 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[6] ---loss: 13.51436 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[7] ---loss: 17.23384 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[8] ---loss: 13.62076 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[9] ---loss: 9.69718 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[10] ---loss: 10.03317 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[11] ---loss: 6.70868 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[12] ---loss: 9.66001 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[13] ---loss: 9.46227 ---time: 0.07 ---prepare time : 0.00\n",
      "iter[14] ---loss: 5.34758 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[15] ---loss: 6.51006 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[16] ---loss: 4.68934 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[17] ---loss: 3.77096 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[18] ---loss: 8.08420 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[19] ---loss: 7.66438 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[20] ---loss: 6.91215 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[21] ---loss: 6.78875 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[22] ---loss: 3.32371 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[23] ---loss: 6.48403 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[24] ---loss: 7.49427 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[25] ---loss: 6.53624 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[26] ---loss: 4.83012 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[27] ---loss: 4.19424 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[28] ---loss: 2.34896 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[29] ---loss: 6.07780 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[30] ---loss: 8.97008 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[31] ---loss: 6.25283 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[32] ---loss: 7.29840 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[33] ---loss: 3.43178 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[34] ---loss: 4.63041 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[35] ---loss: 3.38511 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[36] ---loss: 4.44353 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[37] ---loss: 2.91535 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[38] ---loss: 3.93365 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[39] ---loss: 2.00000 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[40] ---loss: 6.95946 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[41] ---loss: 6.02299 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[42] ---loss: 0.79869 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[43] ---loss: 1.08422 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[44] ---loss: 1.80949 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[45] ---loss: 3.46684 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[46] ---loss: 2.00000 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[47] ---loss: 5.72782 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[48] ---loss: 3.15098 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[49] ---loss: 2.50977 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[50] ---loss: 6.46975 ---time: 0.07 ---prepare time : 0.00\n",
      "iter[51] ---loss: 1.59671 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[52] ---loss: 2.17224 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[53] ---loss: 4.22796 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[54] ---loss: 1.31715 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[55] ---loss: 2.92452 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[56] ---loss: 4.16008 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[57] ---loss: 0.51911 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[58] ---loss: 1.64726 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[59] ---loss: 3.66942 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[60] ---loss: 2.99186 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[61] ---loss: 2.93452 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[62] ---loss: 3.94430 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[63] ---loss: 4.39709 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[64] ---loss: 4.64331 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[65] ---loss: 5.22539 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[66] ---loss: 2.17089 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[67] ---loss: 2.70279 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[68] ---loss: 5.15289 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[69] ---loss: 1.12384 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[70] ---loss: 1.00000 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[71] ---loss: 1.22136 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[72] ---loss: 1.12609 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[73] ---loss: 1.23483 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[74] ---loss: 2.02491 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[75] ---loss: 1.21017 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[76] ---loss: 1.98391 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[77] ---loss: 1.06912 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[78] ---loss: 1.35396 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[79] ---loss: 2.85038 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[80] ---loss: 3.14141 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[81] ---loss: 0.82072 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[82] ---loss: 2.14131 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[83] ---loss: 3.94020 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[84] ---loss: 4.14959 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[85] ---loss: 0.08932 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[86] ---loss: 1.00000 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[87] ---loss: 1.93403 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[88] ---loss: 4.24316 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[89] ---loss: 1.61625 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[90] ---loss: 2.01306 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[91] ---loss: 2.66621 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[92] ---loss: 1.79519 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[93] ---loss: 1.09169 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[94] ---loss: 1.82936 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[95] ---loss: 0.97541 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[96] ---loss: 5.22671 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[97] ---loss: 3.10549 ---time: 0.05 ---prepare time : 0.00\n",
      "iter[98] ---loss: 1.23035 ---time: 0.06 ---prepare time : 0.00\n",
      "iter[99] ---loss: 2.11402 ---time: 0.06 ---prepare time : 0.00\n",
      "Model saved at ./save/TransE_99.ckpt\n"
     ]
    }
   ],
   "source": [
    "# 训练模型\n",
    "args.max_iter = 100\n",
    "\n",
    "config = tf.ConfigProto()\n",
    "config.gpu_options.allow_growth = False\n",
    "config.log_device_placement = False\n",
    "config.allow_soft_placement = True\n",
    "config.gpu_options.per_process_gpu_memory_fraction=0.68\n",
    "session = tf.Session(config=config)\n",
    "session.as_default()\n",
    "\n",
    "tf.initialize_all_variables().run(session=session)\n",
    "saver = tf.train.Saver()\n",
    "norm_rel = session.run(tf.nn.l2_normalize(model.embedding_relation, dim =1))\n",
    "session.run(tf.assign(model.embedding_relation, norm_rel))\n",
    "norm_ent = session.run(tf.nn.l2_normalize(model.embedding_entity, dim =1))\n",
    "session.run(tf.assign(model.embedding_entity, norm_ent))\n",
    "\n",
    "test_input_queue = JoinableQueue()\n",
    "test_output_queue = Queue()\n",
    "\n",
    "workers = list()\n",
    "for i in range(args.n_worker):\n",
    "    worker = Process(target=test_job, args=(test_input_queue, test_output_queue, model.hr_t, model.tr_h ))\n",
    "    worker.start()\n",
    "    workers.append(worker)\n",
    "\n",
    "\n",
    "for n_iter in range(args.max_iter):\n",
    "    accu_loss =0.\n",
    "    batch = 0\n",
    "    num_batch = model.num_triple_train/args.batch_size\n",
    "    start_time = timeit.default_timer()\n",
    "    prepare_time = 0.\n",
    "\n",
    "    for tp, tn , t in  model.training_data_batch(batch_size= args.batch_size):\n",
    "        l, _, l_every, norm_e = session.run([loss, op_train, loss_every, norm_entity], {train_triple_positive_input:tp, train_triple_negative_input: tn})\n",
    "        accu_loss += l\n",
    "        batch += 1\n",
    "        print('[%.2f sec](%d/%d): -- loss: %.5f' %(timeit.default_timer()-start_time, batch, num_batch , l), end='\\r')\n",
    "        prepare_time += t\n",
    "    print('iter[%d] ---loss: %.5f ---time: %.2f ---prepare time : %.2f' %(n_iter, accu_loss, timeit.default_timer()-start_time, prepare_time))\n",
    "\n",
    "    if n_iter % args.save_per == 0 or n_iter ==0 or n_iter == args.max_iter-1:\n",
    "        save_path = saver.save(session, os.path.join('./save/TransE_' + str(n_iter) + '.ckpt'))\n",
    "        print('Model saved at %s' % save_path)\n",
    "\n",
    "    if n_iter %args.evaluate_per_iteration == 0 or n_iter ==0 or n_iter == args.max_iter-1:\n",
    "        rank_head = []\n",
    "        rank_tail = []\n",
    "        filter_rank_head = []\n",
    "        filter_rank_tail = []\n",
    "\n",
    "        norm_rank_head = []\n",
    "        norm_rank_tail = []\n",
    "        norm_filter_rank_head = []\n",
    "        norm_filter_rank_tail = []\n",
    "\n",
    "        start = timeit.default_timer()\n",
    "        testing_data = model.testing_data\n",
    "        hr_t = model.hr_t\n",
    "        tr_h = model.tr_h\n",
    "        n_test = args.n_test\n",
    "        if n_iter == args.max_iter-1:\tn_test = model.num_triple_test      \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 测试一个样本函数\n",
    "def test_one_sample(model, trp, session):\n",
    "    t = trp\n",
    "    id_replace_head , id_replace_tail, norm_id_replace_head , norm_id_replace_tail  = session.run([head_rank, tail_rank, norm_head_rank, norm_tail_rank], {test_triple:t})\n",
    "    hr_t = model.hr_t\n",
    "    tr_h = model.tr_h\n",
    "    \n",
    "    hrank = 0\n",
    "    fhrank = 0\n",
    "    predicted_head_tmp = []\n",
    "    for i in range(len(id_replace_head)):\n",
    "        val = id_replace_head[-i-1]\n",
    "        predicted_head_tmp.append(val)\n",
    "        if val == t[0]:\n",
    "            break\n",
    "        else: \n",
    "            hrank += 1\n",
    "            fhrank += 1 \n",
    "            if val in tr_h[(t[2],t[1])]:\n",
    "                fhrank -= 1\n",
    "    predicted_head_tmp = [id_replace_head[-i-1] for i in range(len(id_replace_head))]\n",
    "    \n",
    "    norm_hrank = 0\n",
    "    norm_fhrank = 0\n",
    "    norm_predicted_head_tmp = []\n",
    "    for i in range(len(norm_id_replace_head)):\n",
    "        val = norm_id_replace_head[-i-1]\n",
    "        norm_predicted_head_tmp.append(val)\n",
    "        if val == t[0]:\n",
    "            break\n",
    "        else: \n",
    "            norm_hrank += 1\n",
    "            norm_fhrank += 1 \n",
    "            if val in tr_h[(t[2],t[1])]:\n",
    "                norm_fhrank -= 1\n",
    "    norm_predicted_head_tmp = [id_replace_head[-i-1] for i in range(len(norm_id_replace_head))]\n",
    "\n",
    "    trank = 0\n",
    "    ftrank = 0\n",
    "    predicted_tail_tmp = []\n",
    "    for i in range(len(id_replace_tail)):\n",
    "        val = id_replace_tail[-i-1]\n",
    "        predicted_tail_tmp.append(val)\n",
    "        if val == t[2]:\n",
    "            break\n",
    "        else:\n",
    "            trank += 1\n",
    "            ftrank += 1\n",
    "            if val in hr_t[(t[0], t[1])]:\n",
    "                ftrank -= 1\n",
    "    predicted_tail_tmp = [id_replace_tail[-i-1] for i in range(len(id_replace_tail))]\n",
    "\n",
    "    norm_trank = 0\n",
    "    norm_ftrank = 0\n",
    "    norm_predicted_tail_tmp = []\n",
    "    for i in range(len(norm_id_replace_tail)):\n",
    "        val = norm_id_replace_tail[-i-1]\n",
    "        norm_predicted_tail_tmp.append(val)\n",
    "        if val == t[2]:\n",
    "            break\n",
    "        else:\n",
    "            norm_trank += 1\n",
    "            norm_ftrank += 1\n",
    "            if val in hr_t[(t[0], t[1])]:\n",
    "                norm_ftrank -= 1\n",
    "    norm_predicted_tail_tmp = [id_replace_tail[-i-1] for i in range(len(norm_id_replace_tail))]\n",
    "    \n",
    "    return hrank, fhrank, trank, ftrank, norm_hrank, norm_fhrank, norm_trank, norm_ftrank, \\\n",
    "            predicted_head_tmp, predicted_tail_tmp, norm_predicted_head_tmp, norm_predicted_tail_tmp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.00 sec] --- testing[1/4]\n",
      "[0.13 sec] --- testing[2/4]\n",
      "[0.13 sec] --- testing[3/4]\n",
      "[0.13 sec] --- testing[4/4]\n",
      "iter:99 --mean rank: 1.75 --hit@10: 1.00\n",
      "iter:99 --filter mean rank: 1.38 --filter hit@10: 1.00\n",
      "iter:99 --norm mean rank: 1.62 --norm hit@10: 1.00\n",
      "iter:99 --norm filter mean rank: 1.25 --norm filter hit@10: 1.00\n"
     ]
    }
   ],
   "source": [
    "# 测试模型\n",
    "predicted_tail = []\n",
    "norm_predicted_tail = []\n",
    "predicted_head = []\n",
    "norm_predicted_head = []\n",
    "\n",
    "rank_head = []\n",
    "rank_tail = []\n",
    "filter_rank_head = []\n",
    "filter_rank_tail = []\n",
    "\n",
    "norm_rank_head = []\n",
    "norm_rank_tail = []\n",
    "norm_filter_rank_head = []\n",
    "norm_filter_rank_tail = []\n",
    "\n",
    "start = timeit.default_timer()\n",
    "testing_data = model.testing_data\n",
    "# hr_t = model.hr_t\n",
    "# tr_h = model.tr_h\n",
    "n_test = args.n_test\n",
    "if n_iter == args.max_iter-1:\tn_test = model.num_triple_test\n",
    "predicted_tail = []\n",
    "norm_predicted_tail = []\n",
    "predicted_head = []\n",
    "norm_predicted_head = []\n",
    "for i in range(n_test):\n",
    "    print('[%.2f sec] --- testing[%d/%d]' %(timeit.default_timer()-start, i+1, n_test))\n",
    "    t = testing_data[i]\n",
    "    hrank, fhrank, trank, ftrank, norm_hrank, norm_fhrank, norm_trank, norm_ftrank, \\\n",
    "            predicted_head_tmp, predicted_tail_tmp, norm_predicted_head_tmp, norm_predicted_tail_tmp = test_one_sample(model, t, session)\n",
    "#     print(hrank, fhrank, trank, ftrank, norm_hrank, norm_fhrank, norm_trank, norm_ftrank)\n",
    "    rank_head.append(hrank)\n",
    "    rank_tail.append(trank)\n",
    "    filter_rank_head.append(fhrank)\n",
    "    filter_rank_tail.append(ftrank)\n",
    "\n",
    "    norm_rank_head.append(norm_hrank)\n",
    "    norm_rank_tail.append(norm_trank)\n",
    "    norm_filter_rank_head.append(norm_fhrank)\n",
    "    norm_filter_rank_tail.append(norm_ftrank)\n",
    "\n",
    "    predicted_tail.append(predicted_tail_tmp)\n",
    "    norm_predicted_tail.append(norm_predicted_tail_tmp)\n",
    "    predicted_head.append(predicted_head_tmp)\n",
    "    norm_predicted_head.append(norm_predicted_head_tmp)\n",
    "\n",
    "mean_rank_head = np.sum(rank_head, dtype=np.float32)/n_test\n",
    "mean_rank_tail = np.sum(rank_tail, dtype=np.float32)/n_test\n",
    "# print('rank_tail:', rank_tail)\n",
    "# # print('predicted_tail:', predicted_tail)\n",
    "# print('norm_predicted_tail:', norm_predicted_tail)\n",
    "# print('predicted_head:', predicted_head)\n",
    "# print('norm_predicted_head:', norm_predicted_head)\n",
    "\n",
    "filter_mean_rank_head = np.sum(filter_rank_head, dtype=np.float32)/n_test\n",
    "filter_mean_rank_tail = np.sum(filter_rank_tail, dtype=np.float32)/n_test\n",
    "\n",
    "norm_mean_rank_head = np.sum(norm_rank_head, dtype=np.float32)/n_test\n",
    "norm_mean_rank_tail = np.sum(norm_rank_tail, dtype=np.float32)/n_test\n",
    "norm_filter_mean_rank_head = np.sum(norm_filter_rank_head, dtype=np.float32)/n_test\n",
    "norm_filter_mean_rank_tail = np.sum(norm_filter_rank_tail, dtype=np.float32)/n_test\n",
    "\n",
    "hit10_head = np.sum(np.asarray(np.asarray(rank_head)<10 , dtype=np.float32))/n_test\n",
    "hit10_tail = np.sum(np.asarray(np.asarray(rank_tail)<10 , dtype=np.float32))/n_test\n",
    "filter_hit10_head = np.sum(np.asarray(np.asarray(filter_rank_head)<10 , dtype=np.float32))/n_test\n",
    "filter_hit10_tail = np.sum(np.asarray(np.asarray(filter_rank_tail)<10 , dtype=np.float32))/n_test\n",
    "\n",
    "norm_hit10_head = np.sum(np.asarray(np.asarray(norm_rank_head)<10 , dtype=np.float32))/n_test\n",
    "norm_hit10_tail = np.sum(np.asarray(np.asarray(norm_rank_tail)<10 , dtype=np.float32))/n_test\n",
    "norm_filter_hit10_head = np.sum(np.asarray(np.asarray(norm_filter_rank_head)<10 , dtype=np.float32))/n_test\n",
    "norm_filter_hit10_tail = np.sum(np.asarray(np.asarray(norm_filter_rank_tail)<10 , dtype=np.float32))/n_test\n",
    "\n",
    "print('iter:%d --mean rank: %.2f --hit@10: %.2f' %(n_iter, (mean_rank_head+ mean_rank_tail)/2, (hit10_tail+hit10_head)/2))\n",
    "print('iter:%d --filter mean rank: %.2f --filter hit@10: %.2f' %(n_iter, (filter_mean_rank_head+ filter_mean_rank_tail)/2, (filter_hit10_tail+filter_hit10_head)/2))\n",
    "\n",
    "print('iter:%d --norm mean rank: %.2f --norm hit@10: %.2f' %(n_iter, (norm_mean_rank_head+ norm_mean_rank_tail)/2, (norm_hit10_tail+norm_hit10_head)/2))\n",
    "print('iter:%d --norm filter mean rank: %.2f --norm filter hit@10: %.2f' %(n_iter, (norm_filter_mean_rank_head+ norm_filter_mean_rank_tail)/2, (norm_filter_hit10_tail+norm_filter_hit10_head)/2))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['北京市', '包含区', '海淀区'] [15, 5, 9]\n",
      "-- trank: 0 ['海淀区', '北京市', '清华大学']\n",
      "-- norm_trank 0 ['海淀区', '北京市', '清华大学']\n",
      "-- hrank 0 ['北京市', '海淀区', '北京大学']\n",
      "-- norm_hrank 0 ['北京市', '海淀区', '北京大学']\n"
     ]
    }
   ],
   "source": [
    "# 测试一个样本,、\n",
    "# 下面我们选取单个样本进行测试，\n",
    "# 输出tail以及head prediction的rank，并出给top k的list，\n",
    "# 这里k取值长度为rank+1+3，方便大家观察排序为正确实体前后的实体分别是什么。\n",
    "\n",
    "sample = ['北京市','包含区','海淀区']\n",
    "# sample = ['鼓楼区','位于省','江苏省']\n",
    "# sample = ['合肥市','位于省','安徽省']\n",
    "# sample = ['浙江大学','位于市','杭州市']\n",
    "trp = [model.entity2id[sample[0]],\n",
    "       model.relation2id[sample[1]],\n",
    "       model.entity2id[sample[2]]]\n",
    "print(sample, trp)\n",
    "\n",
    "hrank, fhrank, trank, ftrank, norm_hrank, norm_fhrank, norm_trank, norm_ftrank, \\\n",
    "            predicted_head_tmp, predicted_tail_tmp, norm_predicted_head_tmp, norm_predicted_tail_tmp = test_one_sample(model, trp, session)\n",
    "\n",
    "tail_prediction = [model.id2entity[ent] for ent in predicted_tail_tmp]\n",
    "head_prediction = [model.id2entity[ent] for ent in predicted_head_tmp]\n",
    "\n",
    "norm_tail_prediction = [model.id2entity[ent] for ent in norm_predicted_tail_tmp]\n",
    "norm_head_prediction = [model.id2entity[ent] for ent in norm_predicted_head_tmp]\n",
    "\n",
    "print('-- trank:', trank, tail_prediction[:trank +3])\n",
    "print('-- norm_trank', norm_trank, norm_tail_prediction[:norm_trank +3])\n",
    "print('-- hrank', hrank, head_prediction[:hrank +3])\n",
    "print('-- norm_hrank', norm_hrank, norm_head_prediction[:norm_hrank +3])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FINISHED~\n"
     ]
    }
   ],
   "source": [
    "# kill all the worker process\n",
    "num_worker = 0\n",
    "for p in workers:\n",
    "    num_worker += 1\n",
    "    p.terminate()\n",
    "    #print('kill worker %d'%num_worker)\n",
    "session.close()\n",
    "print(\"FINISHED~\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "本demo中不包括调参的部分，有兴趣的同学可以自行尝试不同的参数组合，并观察对模型训练和预测结果的影响 :-)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
